package com.sz.szaiagent.rag;

import com.alibaba.cloud.ai.dashscope.api.DashScopeApi;
import jakarta.annotation.Resource;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;

import java.util.List;

import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgDistanceType.COSINE_DISTANCE;
import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIndexType.HNSW;

/**
 * PostgreSQL向量存储配置类
 * 该类用于配置基于PostgreSQL的向量存储，结合DashScope嵌入模型，提升文档的存储和检索效率。
 */
@Configuration
public class PgVectorVectorStoreConfig {

    @Resource
    private LoveAppDocumentLoader loveAppDocumentLoader;

    /**
     * 创建PostgreSQL向量存储实例
     * @param jdbcTemplate JDBC模板，用于数据库操作
     * @param dashscopeEmbeddingModel DashScope嵌入模型，用于文档向量化
     * @return 配置好的PostgreSQL向量存储实例
     */
    @Bean
    public VectorStore pgVectorVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel dashscopeEmbeddingModel) {
        VectorStore vectorStore = PgVectorStore.builder(jdbcTemplate, dashscopeEmbeddingModel)
                .dimensions(1536)                    // 设置向量维度
                .distanceType(COSINE_DISTANCE)       // 设置距离类型
                .indexType(HNSW)                     // 设置索引类型
                .initializeSchema(true)              // 初始化数据库模式
                .schemaName("public")                // 设置模式名称
                .vectorTableName("vector_store")     // 设置向量表名称
                .maxDocumentBatchSize(10000)         // 设置最大文档批处理大小
                .build(); // 构建PostgreSQL向量存储实例
        // 加载文档
        List<Document> documents = loveAppDocumentLoader.loadMarkdowns(); // 加载Markdown文档
        vectorStore.add(documents); // 将文档添加到向量存储
        return vectorStore; // 返回配置好的向量存储实例
    }
}

